/*
 * ModelAssess.h
 *
 *  Created on: May 3, 2012
 *      Author: Simon Whelan
 */

#ifndef MODELASSESS_H_
#define MODELASSESS_H_
#ifdef ISUNIX
#include "optimise.h"
#include "TreeList.h"
#include "model.h"
#include "interface.h"
#else
#include "../PhyloLib/optimise.h"
#include "../PhyloLib/TreeList.h"
#include "../PhyloLib/model.h"
#include "../PhyloLib/interface.h"
#endif

struct s_sequences {
	string Name;
	string Sequence;
};

vector <s_sequences> ReadSequences(string seqin);								// Function reading a set of FASTA sequences to a set of strings

#define APPROX_STARTSIZE 15			// Starting size of the approximation in the Forward algorithm
#define DEFAULT_APPROX_ACC 1.0E-5	// Accuracy of the approximate Forward algorithm likelihood
#define MAX_TIME 1.5				// Maximum time in GuessParameters
class CPairHMM : public CBaseModel {
public:
	// Constructor function
	CPairHMM(s_sequences SeqA, s_sequences SeqB, string ABET);
	// Destructor function
	~CPairHMM();
	// Basic IO functions
	bool CreateModel(EModel ModChoice, double IndelRate);						// Currently a semi-blank function that creates m_vdMatModel and m_dIndel.
	// Primary (immutable) variables
	int m_iChar;						// Size of characters
	string m_sNameA;					// Name of sequence A
	string m_sNameB;					// Name of sequence B
	vector <int> m_viSeqA;				// Sequence A stored as integers mapping to m_sABET
	vector <int> m_viSeqB;				// Sequence B stored as integers mapping to m_s_ABET
	string m_sABET;						// Alphabet of characters. Needed for prettiness
	vector <double> m_vdQMatModel;		// The match probability matrix. Exactly equivalent to a P(t) matrix from phylogenetics
	vector <double> m_vdDataFreq;		// Frequencies taken from the data
	CBaseModel *m_pModel;						// The evolutionary model used
	CTree *m_pDummyTree;						// The dummy tree used in the model
	CData *m_pDummyData;						// The dummy data used in the model
	// Secondary (mutable) variables
	CQPar *m_dTime;								// The evolutionary distance in the model
	CQPar *m_dIndelRate;						// The parameter controlling indel rate
	CQPar *m_dFractionInsert;					// The fraction of indels that are insertions
	CQPar *m_dProbExtend;						// The indel extension probability. 1/ProbExtend is the mean gap length
	vector <double> m_vdPTModel;				// The match probability matrix. Exactly equivalent to a P(t) matrix from phylogenetics
	vector <double> m_vdStateFreq;			// The state frequencies
	double m_dProbIndel;						// The probability of an indel

	// Transition probabilities between hidden states
	vector <double> m_vdMatchTo, m_vdInsTo, m_vdDelTo;		// [0] = To Match; [1] To Insert; [2] to Delete

	// Likelihood stores
	double m_dViterbi_lnL;						// Log probability of the Viterbi path
	double m_dForward_lnL;						// Log probability of the Forward algorithm
	vector <vector <int> > m_viAln;				// The alignment returned by the Viterbi algorithm
	bool m_bDoApprox;							// Whether the approximate version of the Forward algorithm will be used
	int m_iApproxSize;							// The size around the previous best entry of the previous line in the DP matrix. Used when forming m_viApproxPathSize;
	vector <int> m_viApproxPath;				// Takes the viterbi path and uses it as the guide to approximate the likelihood. Stores the location of the second sequence (j-th coordinate)
	vector <int> m_viApproxPathSize;			// The space around m_viApproxPath that's stored for likelihood calculations

	// Calculation storage spaces. These storage spaces are expected to be empty on calls to the associated algorithms
	// Viterbi space. vector <double> holds log-probabilities
	vector <double> m_vdMatchMat;				// The match matrix probability store for calculations
	vector <double> m_vdInsMat;					// The insert (SeqA only) probability store for calculations
	vector <double> m_vdDelMat;					// The delete (SeqB only) probability store for calculations
	// Forward space. vector <CProb> holds probabilities
	CProb *m_pProbMatchMat;						// The match matrix probability store for calculations
	CProb *m_pProbInsMat;						// The insert (SeqA only) probability store for calculations
	CProb *m_pProbDelMat;						// The delete (SeqB only) probability store for calculations

	// Statistical performance measures
	void OutputTimeSurface(ostream &os = cout,double low = 0.0, double high = 1.0);		// Output the lnL surface for t
	void OutputIndelRateSurface(ostream &os = cout,double low = 0.0, double high = 1.0);	// Output the lnL surface for the indel rate
	void OutputFractInsSurface(ostream &os = cout,double low = 0.0, double high = 1.0);	// Output the lnL surface for the fraction insert
	void OutputProbExtendSurface(ostream &os = cout,double low = 0.0, double high = 1.0);	// Output the lnL surface for the probability of extend

	// Primary function calls
	bool BuildModel(bool DoLog);				// Builds the model;
	void GuessParameters();						// Builds a viterbi alignment and guesses the parameters. Currently very basic
	bool OutModel(ostream &os = cout);			// Output the model to outstream
	double Viterbi(bool StorePath = false);		// Calculates the optimal alignment path given the data and model (standard pairHMM viterbi algorithm with a few small alterations) and returns its log-likelihood.
												// If(StorePath) it goes into m_iApproxPath.
	double Forward();							// Calculates the likelihood of two sequences observed conditional on an indel model (forward-algorithm for probability)

	double DoAln(ostream &os = cout);			// Output an alignment. If no alignment exists in m_viAln, then it will build one
	// Functions associated with approximate calculations
	double ApproxForward(double Accuracy = DEFAULT_APPROX_ACC);	// Approximate version of the forward algorithm
	bool BuildApproxPath(double Accuracy,bool reset=true);		// Builds the approximate path for the forward algorithm. Basically builds a viterbi path and stores it in m_vdApproxPath
	bool ForceAccuracy(double Accuracy);						// Chooses the value of m_iApproxSize that ensures Accuracy degree of accuracy in the likelihood function
	void CleanApprox();											// Clean the approximation
	// Secondary function calls
	// Functions getting transition probabilities for Viterbi
	// 1. Going to match; always from (i-1,j-1)
	double Match2Match(int i, int j);
	double Ins2Match(int i, int j);
	double Del2Match(int i, int j);
	// 2. Going to Ins; always from (i-1,j)
	double Match2Ins(int i, int j);
	double Ins2Ins(int i, int j);
	double Del2Ins(int i, int j);
	// 3. Going to Del; always from (i,j-1)
	double Match2Del(int i, int j);
	double Del2Del(int i, int j);
	double Ins2Del(int i, int j);  // Special state to allow insert to delete state
	// Same functions getting transition probabilities for Forward
	// 1. Going to match; always from (i-1,j-1)
	CProb &PrMatch2Match(int i, int j);
	CProb &PrIns2Match(int i, int j);
	CProb &PrDel2Match(int i, int j);
	// 2. Going to Ins; always from (i-1,j)
	CProb &PrMatch2Ins(int i, int j);
	CProb &PrIns2Ins(int i, int j);
	CProb &PrDel2Ins(int i, int j);
	// 3. Going to Del; always from (i,j-1)
	CProb &PrMatch2Del(int i, int j);
	CProb &PrDel2Del(int i, int j);
	CProb &PrIns2Del(int i, int j);  // Special state to allow insert to delete state

	// Stuff that's overwritten from CBaseModel
	bool m_bDoForwardlnL;
	double lnL(bool BLANKTHING = true) { if(m_bDoForwardlnL) { return Forward(); } else { return Viterbi(); } };
};

double DumbOptimiser(CPairHMM *Model);		// Function that sequentially estimates all the parameters in teh model


#endif /* MODELASSESS_H_ */
